{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "saving_loading_dqn.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Stable-Baselines-Team/rl-colab-notebooks/blob/master/saving_loading_dqn.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyyN-2qyK_T2",
        "colab_type": "text"
      },
      "source": [
        "# Stable Baselines, a Fork of OpenAI Baselines - Training, Saving and Loading\n",
        "\n",
        "Github Repo: [https://github.com/hill-a/stable-baselines](https://github.com/hill-a/stable-baselines)\n",
        "\n",
        "Medium article: [https://medium.com/@araffin/stable-baselines-a-fork-of-openai-baselines-df87c4b2fc82](https://medium.com/@araffin/stable-baselines-a-fork-of-openai-baselines-df87c4b2fc82)\n",
        "\n",
        "[RL Baselines Zoo](https://github.com/araffin/rl-baselines-zoo) is a collection of pre-trained Reinforcement Learning agents using Stable-Baselines.\n",
        "\n",
        "It also provides basic scripts for training, evaluating agents, tuning hyperparameters and recording videos.\n",
        "\n",
        "Documentation is available online: [https://stable-baselines.readthedocs.io/](https://stable-baselines.readthedocs.io/)\n",
        "\n",
        "## Install Dependencies and Stable Baselines Using Pip\n",
        "\n",
        "List of full dependencies can be found in the [README](https://github.com/hill-a/stable-baselines).\n",
        "\n",
        "```\n",
        "sudo apt-get update && sudo apt-get install cmake libopenmpi-dev zlib1g-dev\n",
        "```\n",
        "\n",
        "\n",
        "```\n",
        "pip install stable-baselines[mpi]\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gWskDE2c9WoN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Stable Baselines only supports tensorflow 1.x for now\n",
        "%tensorflow_version 1.x\n",
        "!apt install swig cmake libopenmpi-dev zlib1g-dev\n",
        "!pip install stable-baselines[mpi]==2.10.2 box2d box2d-kengz\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FtY8FhliLsGm",
        "colab_type": "text"
      },
      "source": [
        "## Import policy, RL agent, ..."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BIedd7Pz9sOs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gym\n",
        "import numpy as np\n",
        "\n",
        "from stable_baselines import DQN"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RapkYvTXL7Cd",
        "colab_type": "text"
      },
      "source": [
        "## Create the Gym env and instantiate the agent\n",
        "\n",
        "For this example, we will use Lunar Lander environment.\n",
        "\n",
        "\"Landing outside landing pad is possible. Fuel is infinite, so an agent can learn to fly and then land on its first attempt. Four discrete actions available: do nothing, fire left orientation engine, fire main engine, fire right orientation engine. \"\n",
        "\n",
        "Lunar Lander environment: [https://gym.openai.com/envs/LunarLander-v2/](https://gym.openai.com/envs/LunarLander-v2/)\n",
        "\n",
        "![Lunar Lander](https://cdn-images-1.medium.com/max/960/1*f4VZPKOI0PYNWiwt0la0Rg.gif)\n",
        "\n",
        "Note: vectorized environments allow to easily multiprocess training. In this example, we are using only one process, hence the DummyVecEnv.\n",
        "\n",
        "We chose the MlpPolicy because input of CartPole is a feature vector, not images.\n",
        "\n",
        "The type of action to use (discrete/continuous) will be automatically deduced from the environment action space\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pUWGZp3i9wyf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env = gym.make('LunarLander-v2')\n",
        "\n",
        "\n",
        "model = DQN('MlpPolicy', env, learning_rate=1e-3, prioritized_replay=True, verbose=1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4efFdrQ7MBvl",
        "colab_type": "text"
      },
      "source": [
        "We create a helper function to evaluate the agent:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "63M8mSKR-6Zt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def evaluate(model, num_steps=1000):\n",
        "  \"\"\"\n",
        "  Evaluate a RL agent\n",
        "  :param model: (BaseRLModel object) the RL Agent\n",
        "  :param num_steps: (int) number of timesteps to evaluate it\n",
        "  :return: (float) Mean reward for the last 100 episodes\n",
        "  \"\"\"\n",
        "  episode_rewards = [0.0]\n",
        "  obs = env.reset()\n",
        "  for i in range(num_steps):\n",
        "      # _states are only useful when using LSTM policies\n",
        "      action, _states = model.predict(obs)\n",
        "\n",
        "      obs, reward, done, info = env.step(action)\n",
        "      \n",
        "      # Stats\n",
        "      episode_rewards[-1] += reward\n",
        "      if done:\n",
        "          obs = env.reset()\n",
        "          episode_rewards.append(0.0)\n",
        "  # Compute mean reward for the last 100 episodes\n",
        "  mean_100ep_reward = round(np.mean(episode_rewards[-100:]), 1)\n",
        "  print(\"Mean reward:\", mean_100ep_reward, \"Num episodes:\", len(episode_rewards))\n",
        "  \n",
        "  return mean_100ep_reward"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zjEVOIY8NVeK",
        "colab_type": "text"
      },
      "source": [
        "Let's evaluate the un-trained agent, this should be a random agent."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xDHLMA6NFk95",
        "colab_type": "code",
        "outputId": "cf55834c-60f8-4255-a80c-b68957890a3c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "# Random Agent, before training\n",
        "mean_reward_before_train = evaluate(model, num_steps=10000)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mean reward: -895.1 Num episodes: 88\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r5UoXTZPNdFE",
        "colab_type": "text"
      },
      "source": [
        "## Train the agent and save it\n",
        "\n",
        "Warning: this may take a while"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e4cfSXIB-pTF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train the agent\n",
        "model.learn(total_timesteps=int(2e4), log_interval=10)\n",
        "# Save the agent\n",
        "model.save(\"dqn_lunar\")\n",
        "del model  # delete trained model to demonstrate loading"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T31dZJYNrJwF",
        "colab_type": "text"
      },
      "source": [
        "## Load the trained agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K1ExgtyZrIA6",
        "colab_type": "code",
        "outputId": "a6805872-3ec3-454d-b00a-9ff2757b403f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "model = DQN.load(\"dqn_lunar\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loading a model without an environment, this model cannot be trained until it has a valid environment.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ygl_gVmV_QP7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Evaluate the trained agent\n",
        "mean_reward = evaluate(model, num_steps=10000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aQDZI5VEGnUq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}
